"""MCP协议核心实现"""

import json
import uuid
from typing import Dict, Any, Optional, List, Callable, Awaitable
from datetime import datetime
from loguru import logger

import sys
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from src.core.context_orchestrator import ContextOrchestrator


class MCPProtocol:
    """MCP协议处理器"""
    
    def __init__(self, orchestrator: Optional[ContextOrchestrator], config: Dict[str, Any]):
        self.orchestrator = orchestrator
        self.config = config
        self.mcp_config = config.get('mcp', {})
        
        # 注册的工具和资源
        self.tools: Dict[str, Callable] = {}
        self.resources: Dict[str, Callable] = {}
        
        # 工具与资源管理器实例
        self.tools_manager = None
        self.resources_manager = None
        
        # 会话管理
        self.sessions: Dict[str, Dict[str, Any]] = {}
        
        # 注册默认工具和资源
        self._register_default_tools()
        self._register_default_resources()
    
    async def initialize(self):
        """初始化MCP协议"""
        try:
            # 如果有orchestrator，初始化它
            if self.orchestrator:
                await self.orchestrator.initialize()
            logger.info("MCP协议初始化完成")
        except Exception as e:
            logger.error(f"MCP协议初始化失败: {e}")
            raise
    
    def _register_default_tools(self):
        """注册默认工具"""
        if self.orchestrator:
            from .tools import MCPTools
            tools_instance = MCPTools(self.orchestrator)
            # 持有工具管理器
            self.tools_manager = tools_instance
            
            # 仍然保留到字典的快捷引用（主要用于回退/兼容）
            for tool_name in ['memory_search', 'knowledge_upload', 'graph_query', 
                              'conversation_create', 'conversation_search', 'create_memory', 
                              'knowledge_search', 'find_related_concepts', 'get_user_context',
                              'system_stats', 'health_check']:
                if tool_name in tools_instance.tools:
                    self.tools[tool_name] = tools_instance.tools[tool_name]['handler']
        else:
            # 如果没有orchestrator，注册一个简单的echo工具
            self.tools['echo'] = self._echo_tool
    
    def _register_default_resources(self):
        """注册默认资源"""
        if self.orchestrator:
            from .resources import MCPResources
            resources_instance = MCPResources(self.orchestrator)
            # 持有资源管理器
            self.resources_manager = resources_instance
            
            # 注册资源处理器（兼容保留）
            for resource_name in ['conversations', 'documents', 'embeddings', 'memories', 
                                 'graph_data', 'user_context', 'system_stats', 'health_status']:
                if resource_name in resources_instance.resources:
                    self.resources[resource_name] = resources_instance.resources[resource_name]['handler']
    
    async def _echo_tool(self, text: str) -> Dict[str, Any]:
        """简单的echo工具，用于测试"""
        return {
            "content": [
                {
                    "type": "text",
                    "text": f"Echo: {text}"
                }
            ]
        }
    
    async def handle_message(self, message: str, session_id: Optional[str] = None) -> str:
        """处理MCP消息"""
        try:
            # 解析JSON消息
            data = json.loads(message)
            
            # 生成会话ID（如果未提供）
            if not session_id:
                session_id = str(uuid.uuid4())
            
            # 初始化会话（如果不存在）
            if session_id not in self.sessions:
                self.sessions[session_id] = {
                    'created_at': datetime.now().isoformat(),
                    'last_activity': datetime.now().isoformat(),
                    'message_count': 0
                }
            
            # 更新会话活动时间
            self.sessions[session_id]['last_activity'] = datetime.now().isoformat()
            self.sessions[session_id]['message_count'] += 1
            
            # 处理不同类型的消息
            if 'method' in data:
                # JSON-RPC请求
                response = await self._handle_jsonrpc_request(data, session_id)
            elif 'id' in data and 'result' in data:
                # JSON-RPC响应
                response = await self._handle_jsonrpc_response(data, session_id)
            elif 'id' in data and 'error' in data:
                # JSON-RPC错误响应
                response = await self._handle_jsonrpc_error(data, session_id)
            else:
                # 未知消息格式
                response = self._create_error_response(
                    None, -32600, "Invalid Request", "未知的消息格式"
                )
            
            return json.dumps(response, ensure_ascii=False)
            
        except json.JSONDecodeError as e:
            logger.error(f"JSON解析错误: {e}")
            return json.dumps(self._create_error_response(
                None, -32700, "Parse error", str(e)
            ), ensure_ascii=False)
        except Exception as e:
            logger.error(f"处理MCP消息失败: {e}")
            return json.dumps(self._create_error_response(
                None, -32603, "Internal error", str(e)
            ), ensure_ascii=False)
    
    async def _handle_jsonrpc_request(self, data: Dict[str, Any], session_id: str) -> Dict[str, Any]:
        """处理JSON-RPC请求"""
        method = data.get('method')
        params = data.get('params', {})
        request_id = data.get('id')
        
        logger.debug(f"处理JSON-RPC请求: {method}")
        
        try:
            # 处理不同的方法
            if method == 'initialize':
                result = await self._handle_initialize(params, session_id)
            elif method == 'tools/list':
                result = await self._handle_tools_list(params, session_id)
            elif method == 'tools/call':
                result = await self._handle_tools_call(params, session_id)
            elif method == 'resources/list':
                result = await self._handle_resources_list(params, session_id)
            elif method == 'resources/read':
                result = await self._handle_resources_read(params, session_id)
            elif method == 'ping':
                result = {'status': 'pong', 'timestamp': datetime.now().isoformat()}
            else:
                return self._create_error_response(
                    request_id, -32601, "Method not found", f"未知方法: {method}"
                )
            
            return self._create_success_response(request_id, result)
            
        except Exception as e:
            logger.error(f"处理方法 {method} 失败: {e}")
            return self._create_error_response(
                request_id, -32603, "Internal error", str(e)
            )
    
    async def _handle_jsonrpc_response(self, data: Dict[str, Any], session_id: str) -> Dict[str, Any]:
        """处理JSON-RPC响应"""
        # 这里可以处理来自客户端的响应
        logger.debug(f"收到JSON-RPC响应: {data.get('id')}")
        return {'status': 'acknowledged'}
    
    async def _handle_jsonrpc_error(self, data: Dict[str, Any], session_id: str) -> Dict[str, Any]:
        """处理JSON-RPC错误响应"""
        logger.error(f"收到JSON-RPC错误: {data.get('error')}")
        return {'status': 'error_acknowledged'}
    
    async def _handle_initialize(self, params: Dict[str, Any], session_id: str) -> Dict[str, Any]:
        """处理初始化请求"""
        client_info = params.get('clientInfo', {})
        capabilities = params.get('capabilities', {})
        
        logger.info(f"客户端初始化: {client_info.get('name', 'Unknown')} v{client_info.get('version', 'Unknown')}")
        
        # 更新会话信息
        self.sessions[session_id].update({
            'client_info': client_info,
            'client_capabilities': capabilities,
            'initialized': True
        })
        
        # 返回服务器信息和能力
        return {
            'protocolVersion': '2024-11-05',
            'serverInfo': {
                'name': 'Memory AI MCP Server',
                'version': '1.0.0'
            },
            'capabilities': {
                'tools': {
                    'listChanged': True
                },
                'resources': {
                    'subscribe': True,
                    'listChanged': True
                },
                'logging': {},
                'experimental': {}
            }
        }
    
    async def _handle_tools_list(self, params: Dict[str, Any], session_id: str) -> Dict[str, Any]:
        """处理工具列表请求"""
        # 优先从工具管理器获取，保证包含完整的输入Schema与描述
        if self.tools_manager is not None:
            tools_list = self.tools_manager.get_tool_list()
            return {'tools': tools_list}
        
        # 回退：基于已注册的处理器构造最小工具信息
        tools_list: List[Dict[str, Any]] = []
        for tool_name, tool_func in self.tools.items():
            safe_desc = (getattr(tool_func, '__doc__', None) or f'{tool_name} 工具')
            tool_info = {
                'name': tool_name,
                'description': safe_desc,
                'inputSchema': {
                    'type': 'object',
                    'properties': {},
                    'required': []
                }
            }
            # 兜底：确保 description 为非空字符串
            desc = tool_info.get('description')
            if desc is None or (isinstance(desc, str) and not desc.strip()):
                tool_info['description'] = f'{tool_name} 工具'
            tools_list.append(tool_info)
        return {'tools': tools_list}
    
    async def _handle_tools_call(self, params: Dict[str, Any], session_id: str) -> Dict[str, Any]:
        """处理工具调用请求"""
        tool_name = params.get('name')
        arguments = params.get('arguments', {})
        
        if self.tools_manager is None and tool_name not in self.tools:
            raise ValueError(f"未知工具: {tool_name}")
        
        # 添加会话信息到参数（透传给工具，供需要者使用）
        arguments = dict(arguments or {})
        arguments['_session_id'] = session_id
        arguments['_user_id'] = self.sessions[session_id].get('user_id', 'anonymous')
        
        if self.tools_manager is not None:
            # 通过工具管理器统一调用，避免参数解包不匹配
            result = await self.tools_manager.call_tool(tool_name, arguments)
            return result
        
        # 回退路径（仅用于echo等简易工具）
        tool_func = self.tools[tool_name]
        # 简易工具沿用原调用方式
        raw = await tool_func(**arguments)
        return {
            'content': [
                {
                    'type': 'text',
                    'text': json.dumps(raw, ensure_ascii=False)
                }
            ]
        }
    
    async def _handle_resources_list(self, params: Dict[str, Any], session_id: str) -> Dict[str, Any]:
        """处理资源列表请求"""
        if self.resources_manager is not None:
            resources_list = self.resources_manager.get_resource_list()
            # 可选：为外部暴露添加 scheme 前缀以保持一致性
            exposed = []
            for r in resources_list:
                exposed.append({
                    'uri': f"memory://{r['uri']}",
                    'name': r['name'],
                    'description': r['description'] or f"{r['uri']} 资源",
                    'mimeType': r['mimeType']
                })
            return {'resources': exposed}
        
        # 回退：基于已注册处理器构造
        resources_list: List[Dict[str, Any]] = []
        for resource_name, resource_func in self.resources.items():
            resource_info = {
                'uri': f'memory://{resource_name}',
                'name': resource_name,
                'description': getattr(resource_func, '__doc__', f'{resource_name} 资源'),
                'mimeType': 'application/json'
            }
            resources_list.append(resource_info)
        return {'resources': resources_list}
    
    async def _handle_resources_read(self, params: Dict[str, Any], session_id: str) -> Dict[str, Any]:
        """处理资源读取请求"""
        uri = params.get('uri', '')
        
        # 解析URI
        if uri.startswith('memory://'):
            resource_name = uri[9:]  # 移除 'memory://' 前缀
        else:
            resource_name = uri
        
        if self.resources_manager is None and resource_name not in self.resources:
            raise ValueError(f"未知资源: {resource_name}")
        
        # 组装传递给资源读取的参数（不直接作为kwargs解包）
        resource_params = dict(params or {})
        resource_params['_session_id'] = session_id
        resource_params['_user_id'] = self.sessions[session_id].get('user_id', 'anonymous')
        
        if self.resources_manager is not None:
            # 使用资源管理器读取，避免参数解包不匹配
            result = await self.resources_manager.read_resource(resource_name, resource_params)
            # 保持返回内容中的URI与请求一致
            if 'contents' in result and result['contents']:
                for item in result['contents']:
                    item['uri'] = uri
            return result
        
        # 回退路径：直接调用处理器（注意：此路径可能因为签名不一致而失败，不建议使用）
        resource_func = self.resources[resource_name]
        raw = await resource_func(**resource_params)
        return {
            'contents': [
                {
                    'uri': uri,
                    'mimeType': 'application/json',
                    'text': json.dumps(raw, ensure_ascii=False)
                }
            ]
        }
    
    def _create_success_response(self, request_id: Any, result: Any) -> Dict[str, Any]:
        """创建成功响应"""
        response = {
            'jsonrpc': '2.0',
            'result': result
        }
        
        if request_id is not None:
            response['id'] = request_id
        
        return response
    
    def _create_error_response(self, request_id: Any, code: int, message: str, data: Any = None) -> Dict[str, Any]:
        """创建错误响应"""
        error = {
            'code': code,
            'message': message
        }
        
        if data is not None:
            error['data'] = data
        
        response = {
            'jsonrpc': '2.0',
            'error': error
        }
        
        if request_id is not None:
            response['id'] = request_id
        
        return response
    
    def register_tool(self, name: str, func: Callable, schema: Optional[Dict[str, Any]] = None):
        """注册工具"""
        self.tools[name] = func
        if schema:
            func._mcp_schema = schema
        logger.info(f"工具已注册: {name}")
    
    def register_resource(self, name: str, func: Callable):
        """注册资源"""
        self.resources[name] = func
        logger.info(f"资源已注册: {name}")
    
    def get_session_info(self, session_id: str) -> Optional[Dict[str, Any]]:
        """获取会话信息"""
        return self.sessions.get(session_id)
    
    def cleanup_inactive_sessions(self, max_inactive_hours: int = 24):
        """清理不活跃的会话"""
        from datetime import timedelta
        
        cutoff_time = datetime.now() - timedelta(hours=max_inactive_hours)
        inactive_sessions = []
        
        for session_id, session_info in self.sessions.items():
            last_activity = datetime.fromisoformat(session_info['last_activity'])
            if last_activity < cutoff_time:
                inactive_sessions.append(session_id)
        
        for session_id in inactive_sessions:
            del self.sessions[session_id]
            logger.info(f"清理不活跃会话: {session_id}")
        
        return len(inactive_sessions)